#include <opencv2/opencv.hpp>

#include "final_exam/pose_solver.h"

void PoseSolver::setCameraParams(const cv::Mat &camera_matrix,
                                 const cv::Mat &dist_coeffs) {
  this->camera_matrix = camera_matrix.clone();
  this->dist_coeffs = dist_coeffs.clone();
}

void PoseSolver::setObjectPoints(
    const std::vector<cv::Point3f> &object_points) {
  this->object_points = object_points;
}

bool PoseSolver::solve(const std::vector<cv::Point2f> &image_points) {
  if (image_points.size() != this->object_points.size() ||
      image_points.size() < 4)
    return false;
  bool success =
      cv::solvePnP(this->object_points, image_points, this->camera_matrix,
                   this->dist_coeffs, this->rvec, this->tvec);
  if (success)
    return true;
  else
    return false;
}

std::vector<cv::Point2f> PoseSolver::project(const cv::Mat &rvec,
                                             const cv::Mat &tvec) {
  std::vector<cv::Point2f> image_points;
  cv::projectPoints(this->object_points, rvec, tvec, this->camera_matrix,
                    this->dist_coeffs, image_points);
  return image_points;
}

cv::Mat PoseSolver::getRvec() const { return this->rvec; }

cv::Mat PoseSolver::getTvec() const { return this->tvec; }